

import os
import sys
import cPickle

import rpyc

from dexen.system import constants 

""" TO BE SET BY LOOPER """
_url = None
_schema_name = None
_pop_name = None
_script_name = None

def ConfigAPIFactory():
    return ConfigAPI(_url, _schema_name, _pop_name, _script_name)

def PopAPIFactory():
    return PopAPI(_url, _schema_name, _pop_name, _script_name)

def _set_attributes(url, schema_name=None, pop_name=None, script_name=None):
    global _url, _schema_name, _pop_name, _script_name
    _url = url
    _schema_name = schema_name
    _pop_name = pop_name
    _script_name = script_name
    
    
class ConfigAPI:       
    def __init__(self, url, schema_name, pop_name, script_name):
        self.schema_name = schema_name
        self.pop_name = pop_name
        self.conn = rpyc.connect(url, port=constants.CONFIG_SRVC_PORT)
        self.conn.root.initialize(schema_name, pop_name)
    
    def setDataStructure(self, attributes):
        self.conn.root.setDataStructure(cPickle.dumps(attributes))

    def setGlobalData(self, data):
        self.conn.root.setGlobalData(cPickle.dumps(data))

    def setObjectives(self, attributes):
        self.conn.root.setObjectives(cPickle.dumps(attributes))
    
    def setImage(self, image_name):
        self.conn.root.setImage(image_name)
    
    def setModel(self, model_name):
        self.conn.root.setModel(model_name)
    
    def setMaxIndividuals(self, max_individuals):
        self.conn.root.setMaxIndividuals(max_individuals)
    
    def setTask(self, script_name, file_name, function_name, 
                run_type=constants.MANUAL, select=None, where=None, count=None, 
                order_by=None, exclusive_use=False):        
        """Register the task in the database by making a remote call to server.
        
        Arguments:
            script_name:
            file_name:
            function_name:
            run_type:
            select:
            where:
            count:
            order_by:
            exclusive_use:
            
        Note:
        exclusive_use is set to False by default.
        
        """
        
        self.conn.root.setScript(script_name, file_name, function_name, 
                                 run_type,select, where, count, order_by, 
                                 exclusive_use)

class PopAPI:
    def __init__(self, url, schema_name, pop_name, script_name):
        self.url = url
        self.schema_name = schema_name
        self.pop_name = pop_name
        self.script_name = script_name
        self.conn = rpyc.connect(url, constants.POP_SRVC_PORT, 
                                 config={"instantiate_custom_exceptions":True})
        self.conn.root.initialize(schema_name, pop_name, script_name)
        self.IMAGE_DIR = os.path.join(constants.OUTPUT_DIR, schema_name, 
                                      pop_name, constants.IMAGE_DIRNAME)
        self.MODEL_DIR = os.path.join(constants.OUTPUT_DIR, schema_name, 
                                      pop_name, constants.MODEL_DIRNAME)

    def _save_image(self, ind):
        filename = "%s.png" %ind["id"]
        f = open(os.path.join(self.IMAGE_DIR, filename), "wb")
        f.write(ind[constants.IMAGE])
        f.close()
        
    def _save_images(self, individuals):
        for ind in individuals:
            if ind.has_key(constants.IMAGE):
                self._save_image(ind)
    
    def setImage(self, ind, image, keep=False):
        ind[constants.IMAGE] = image
        if keep:
            self._save_image(ind)
        
    def getImage(self, ind):
        return ind[constants.IMAGE]

    def killIndividual(self, ind):
        ind[constants.LIFE_STATE] = constants.DEAD

    def reviveIndividual(self, ind):
        ind[constants.LIFE_STATE] = constants.ALIVE
        
    def getGlobalData(self):
        global_data = self.conn.root.getGlobalData()
        return cPickle.loads(global_data)
    
    def download_live_individuals(self, **kwargs):
        inds = self.conn.root.get_live_individuals(**kwargs)
        return cPickle.loads(inds)
    
    def count_live_individuals(self, where=(), at_time=None):
        return self.conn.root.count_live_individuals(where, at_time)
    
    